from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# import pycocotools.coco as coco
# from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import os

import torch.utils.data as data


class COCO(data.Dataset):
    num_classes = 221
    default_resolution = [512, 512]  #:2048*2048
    mean = np.array([0.40789654, 0.44719302, 0.47026115],
                    dtype=np.float32).reshape(1, 1, 3)
    std = np.array([0.28863828, 0.27408164, 0.27809835],
                   dtype=np.float32).reshape(1, 1, 3)

    def __init__(self, opt, split):
        super(COCO, self).__init__()
        self.data_dir = os.path.join("/cfs/DataSets/signs")
        self.img_dir = os.path.join(
            "/cfs/DataSets/signs")
        if split == 'train':
            self.anno_path = os.path.join(
                "../data/annotations/train_test_ids.txt")
        elif split == 'test':
            self.anno_path = os.path.join(
                "../data/annotations/test_ids.txt")
        self.anno_dict = json.loads(open(
            "../data/annotations/annotations.json").read())
        self.max_objs = 10
        self._data_rng = np.random.RandomState(123)
        self._data_rng = np.random.RandomState(123)
        self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],
                                  dtype=np.float32)
        self._eig_vec = np.array([
            [-0.58752847, -0.69563484, 0.41340352],
            [-0.5832747, 0.00994535, -0.81221408],
            [-0.56089297, 0.71832671, 0.41158938]
        ], dtype=np.float32)
        self.class_name = ['i1',
                           'i10',
                           'i11',
                           'i12',
                           'i13',
                           'i14',
                           'i15',
                           'i2',
                           'i3',
                           'i4',
                           'i5',
                           'il100',
                           'il110',
                           'il50',
                           'il60',
                           'il70',
                           'il80',
                           'il90',
                           'io',
                           'ip',
                           'p1',
                           'p10',
                           'p11',
                           'p12',
                           'p13',
                           'p14',
                           'p15',
                           'p16',
                           'p17',
                           'p18',
                           'p19',
                           'p2',
                           'p20',
                           'p21',
                           'p22',
                           'p23',
                           'p24',
                           'p25',
                           'p26',
                           'p27',
                           'p28',
                           'p3',
                           'p4',
                           'p5',
                           'p6',
                           'p7',
                           'p8',
                           'p9',
                           'pa10',
                           'pa12',
                           'pa13',
                           'pa14',
                           'pa8',
                           'pb',
                           'pc',
                           'pg',
                           'ph1.5',
                           'ph2',
                           'ph2.1',
                           'ph2.2',
                           'ph2.4',
                           'ph2.5',
                           'ph2.8',
                           'ph2.9',
                           'ph3',
                           'ph3.2',
                           'ph3.5',
                           'ph3.8',
                           'ph4',
                           'ph4.2',
                           'ph4.3',
                           'ph4.5',
                           'ph4.8',
                           'ph5',
                           'ph5.3',
                           'ph5.5',
                           'pl10',
                           'pl100',
                           'pl110',
                           'pl120',
                           'pl15',
                           'pl20',
                           'pl25',
                           'pl30',
                           'pl35',
                           'pl40',
                           'pl5',
                           'pl50',
                           'pl60',
                           'pl65',
                           'pl70',
                           'pl80',
                           'pl90',
                           'pm10',
                           'pm13',
                           'pm15',
                           'pm1.5',
                           'pm2',
                           'pm20',
                           'pm25',
                           'pm30',
                           'pm35',
                           'pm40',
                           'pm46',
                           'pm5',
                           'pm50',
                           'pm55',
                           'pm8',
                           'pn',
                           'pne',
                           'po',
                           'pr10',
                           'pr100',
                           'pr20',
                           'pr30',
                           'pr40',
                           'pr45',
                           'pr50',
                           'pr60',
                           'pr70',
                           'pr80',
                           'ps',
                           'pw2',
                           'pw2.5',
                           'pw3',
                           'pw3.2',
                           'pw3.5',
                           'pw4',
                           'pw4.2',
                           'pw4.5',
                           'w1',
                           'w10',
                           'w12',
                           'w13',
                           'w16',
                           'w18',
                           'w20',
                           'w21',
                           'w22',
                           'w24',
                           'w28',
                           'w3',
                           'w30',
                           'w31',
                           'w32',
                           'w34',
                           'w35',
                           'w37',
                           'w38',
                           'w41',
                           'w42',
                           'w43',
                           'w44',
                           'w45',
                           'w46',
                           'w47',
                           'w48',
                           'w49',
                           'w5',
                           'w50',
                           'w55',
                           'w56',
                           'w57',
                           'w58',
                           'w59',
                           'w60',
                           'w62',
                           'w63',
                           'w66',
                           'w8',
                           'wo',
                           'i6',
                           'i7',
                           'i8',
                           'i9',
                           'ilx',
                           'p29',
                           'w29',
                           'w33',
                           'w36',
                           'w39',
                           'w4',
                           'w40',
                           'w51',
                           'w52',
                           'w53',
                           'w54',
                           'w6',
                           'w61',
                           'w64',
                           'w65',
                           'w67',
                           'w7',
                           'w9',
                           'pax',
                           'pd',
                           'pe',
                           'phx',
                           'plx',
                           'pmx',
                           'pnl',
                           'prx',
                           'pwx',
                           'w11',
                           'w14',
                           'w15',
                           'w17',
                           'w19',
                           'w2',
                           'w23',
                           'w25',
                           'w26',
                           'w27',
                           'pl0',
                           'pl4',
                           'pl3',
                           'pm2.5',
                           'ph4.4',
                           'pn40',
                           'ph3.3',
                           'ph2.6']

        # self.cat_ids = {v: i for i, v in enumerate(self._valid_ids)}
        # self.voc_color = [(v // 32 * 64 + 64, (v // 8) % 4 * 64, v % 8 * 32) \
        #                  for v in range(1, self.num_classes + 1)]

        # self.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)
        # self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)

        self.split = split
        self.opt = opt
        print('==> initializing traffic signs  {} data.'.format(split))
        # self.anno = json.load(open(self.annot_path))
        # self.coco = coco.COCO(os.path.join(self.data_dir, 'annotations_clip_832_448', 'coco_annotations.json'))
        # image_ids = self.coco.getImgIds()
        self.images_ids = open(self.anno_path).readlines()
        # self.annotations = json.loads(open(self.anno_path).read())
        # self.images = os.listdir(self.img_dir)
        self.num_samples = len(self.images_ids)

        print('Loaded {} {} samples'.format(split, self.num_samples))

    def _to_float(self, x):
        return float("{:.2f}".format(x))

    def convert_eval_format(self, all_bboxes):
        # import pdb; pdb.set_trace()
        detections = []
        for image_id in all_bboxes:
            for cls_ind in all_bboxes[image_id]:
                category_id = 0
                for bbox in all_bboxes[image_id][cls_ind]:
                    bbox[2] -= bbox[0]
                    bbox[3] -= bbox[1]
                    score = bbox[4]
                    bbox_out = list(map(self._to_float, bbox[0:4]))
                    detection = {
                        "image_id": int(image_id),
                        "category_id": int(category_id),
                        "bbox": bbox_out,
                        "score": float("{:.2f}".format(score))
                    }
                    if len(bbox) > 5:
                        extreme_points = list(map(self._to_float, bbox[5:13]))
                        detection["extreme_points"] = extreme_points
                    detections.append(detection)
        return detections

    def __len__(self):
        return self.num_samples

    def save_results(self, results, save_dir):
        json.dump(self.convert_eval_format(results),
                  open('{}/results.json'.format(save_dir), 'w'))

    def run_eval(self, results, save_dir):
        # result_json = os.path.join(save_dir, "results.json")
        # detections  = self.convert_eval_format(results)
        # json.dump(detections, open(result_json, "w"))
        self.save_results(results, save_dir)
        coco_dets = self.coco.loadRes('{}/results.json'.format(save_dir))
        coco_eval = COCOeval(self.coco, coco_dets, "bbox")
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()
